package com.gitee.cirnochat.chatai.client;


import cn.hutool.json.JSONArray;
import com.baomidou.mybatisplus.core.toolkit.IdWorker;
import com.gitee.cirnochat.chatai.model.dto.lke.LkeAskRequest;
import com.gitee.cirnochat.chatai.model.dto.lke.LkeReplyResponse;
import com.gitee.cirnochat.chatai.model.dto.lke.TencentCloudLkeOptions;
import com.gitee.cirnochat.chatai.model.dto.lke.TencentCloudLkeToken;
import com.gitee.cirnochat.chatai.model.entity.AiChatMessage;
import com.gitee.cirnochat.chatai.model.entity.AiChatSession;
import com.gitee.cirnochat.chatai.model.enums.aimessage.ChatRoleType;
import com.gitee.cirnochat.chatai.model.enums.aimessage.MessageType;
import com.gitee.cirnochat.chatai.model.vo.resp.AIChatSessionVo;
import com.gitee.cirnochat.chatai.model.vo.resp.AISteamDoneVo;
import com.gitee.cirnochat.chatai.model.vo.resp.AIStreamMessageVo;
import com.gitee.cirnochat.chatai.model.vo.resp.AiChatSessionMessageVo;
import com.gitee.cirnochat.chatai.service.AiChatMessageService;
import com.gitee.cirnochat.chatai.service.AiChatSessionService;
import com.gitee.cirnochat.chatai.service.AiSummaryService;
import com.gitee.cirnochat.common.model.vo.resp.StreamData;
import com.gitee.cirnochat.common.utils.MarkdownUtil;
import io.socket.client.IO;
import io.socket.client.Socket;
import io.socket.engineio.client.transports.WebSocket;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.json.JSONException;
import org.json.JSONObject;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxSink;

import java.net.URISyntaxException;
import java.time.Duration;
import java.time.LocalDateTime;
import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

@Component
@Slf4j
@RequiredArgsConstructor
public class DeepDuckChatClient {
    private final TencentCloudLkeToken token;

    private final TencentCloudLkeOptions options;

    private final static int MAX_QPS = 2; // 每秒最大请求数

    private final static Semaphore REQUEST = new Semaphore(MAX_QPS);

    private final ThreadPoolTaskExecutor taskExecutor;

    private final AiChatSessionService aiChatSessionService;

    private final AiChatMessageService aiChatMessageService;

    private final AiSummaryService aiSummaryService;



    private Socket socket;

    @PostConstruct
    public void init() {
        this.socket = initClient();
    }

    private Socket initClient() {
        if (socket != null && socket.connected()) {
            log.info("Socket is already connected, no need to reinitialize.");
            return socket;
        }

        String wsUrl = options.getWsUrl();
        String wsCompletionsPath = options.getWsCompletionsPath();

        IO.Options options = IO.Options.builder()
                .setQuery("EIO=4")
                .setPath(wsCompletionsPath)  // path需要以'/'开头
                .setTransports(new String[]{WebSocket.NAME})
                .setAuth(Map.of("token", token.getWsToken()))
                .setTimeout(2000)
                .build();
        // 创建 Socket.IO 客户端实例
        try {
            final Socket socket = IO.socket(wsUrl, options);
            // 监听连接成功事件
            socket.on(Socket.EVENT_CONNECT, args -> log.info("EVENT_CONNECT: {}", Arrays.toString(args)));
            // 监听连接错误事件
            socket.on(Socket.EVENT_CONNECT_ERROR, args -> {
                log.error("EVENT_CONNECT_ERROR: {}", Arrays.toString(args));
                // 如果连接失败，可能需要更新auth，再发起连接
                options.auth.put("token", token.getWsToken());
                socket.disconnect().connect();
            });
            // 监听连接断开事件
            socket.on(Socket.EVENT_DISCONNECT, args ->
                    log.warn("EVENT_DISCONNECT: {}", Arrays.toString(args))
            );
            socket.connect();
            return socket;
        } catch (URISyntaxException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 发送消息到聊天服务
     *
     * @param request 请求对象，包含会话ID和消息内容
     */
    public Flux<StreamData<?>> stream(LkeAskRequest request) {
        try {
            boolean acquired = REQUEST.tryAcquire(1, TimeUnit.MINUTES); // 尝试获取许可，阻塞直到获取到许可或超时
            if (!acquired) {
                throw new RuntimeException("请求许可超时，可能由于QPS限制导致线程阻塞");
            }
            // 请求ID
            String requestId = UUID.randomUUID().toString();
            // 回复开始时间
            LocalDateTime replyStartTime = LocalDateTime.now();
            // 会话ID
            Long sessionId = request.getSessionId() == null ? IdWorker.getId() : request.getSessionId();
            // 用户消息ID
            Long userMessageId = IdWorker.getId();
            // 聊天消息ID
            Long chatMessageId = IdWorker.getId();
            // 封装发送数据
            JSONObject sendPayload = new JSONObject();
            sendPayload.put("request_id", requestId);
            sendPayload.put("content", request.getContent());
            sendPayload.put("stream", "enable");
            sendPayload.put("incremental", true);
            sendPayload.put("session_id", sessionId.toString());
            sendPayload.put("model_name", request.getModel());
            sendPayload.put("search_network",request.getSearchNetwork());
            JSONObject data = new JSONObject();
            data.put("payload", sendPayload);
            return Flux.create(sink -> {
                AtomicReference<String> recordId = new AtomicReference<>(); // 用于存储生成记录的ID
                final StringBuilder sb = new StringBuilder();

                // 异步总结
                Future[] futures = new Future[]{null, null};
                AiChatSession aiChatSession = request.getSessionId() != null ? aiChatSessionService.getById(request.getSessionId()) : null;
                if (aiChatSession == null) {
                    aiChatSession = createAiChatSession(sessionId, request.getUserId());
                    futures[0] = summaryQuestionAsync(sink, aiChatSession, request.getContent());
                }

                // 用户消息更新
                futures[1] = taskExecutor.submit(() -> {
                    AiChatMessage userMessage = new AiChatMessage();
                    userMessage.setId(userMessageId);
                    userMessage.setSessionId(sessionId);
                    userMessage.setContent(request.getContent());
                    userMessage.setModel(request.getModel());
                    userMessage.setRole(ChatRoleType.USER);
                    aiChatMessageService.save(userMessage);
                    sink.next(StreamData.updateMessage(new AiChatSessionMessageVo(userMessage)));
                    return null;
                });

                // 断开/取消时移除监听器
                sink.onCancel(() -> {
                    socket.off("reply");
                    if (recordId.get() != null) {
                        // 停止生成
                        JSONObject payload = new JSONObject();
                        try {
                            JSONObject jsonObject = new JSONObject();
                            jsonObject.put("record_id", recordId.get());
                            payload.put("payload", jsonObject);
                        } catch (JSONException ex) {
                            throw new RuntimeException(ex);
                        }
                        log.info("payload: {}", payload);
                        socket.emit("stop_generation", payload.toString());
                    }
                });

                log.info("发送消息: {}", data);

                socket.emit("send", data)
                        // 监听reply事件
                        .on("reply", args -> {
                            if (args.length > 0) {
                                JSONObject response = (JSONObject) args[0];
                                try {
                                    JSONObject payload = response.getJSONObject("payload");
                                    String recordIdStr = payload.getString("record_id");
                                    boolean isFromSelf = payload.getBoolean("is_from_self");
                                    boolean isFinal = payload.getBoolean("is_final");
                                    String replyContent = payload.getString("content");
                                    if (isFromSelf) {
                                        // 更新消息
                                        aiChatMessageService.updateMessageRecordIdAsync(chatMessageId, recordIdStr);
                                        return;
                                    }
                                    if (recordId.get() == null) {
                                        // 获取生成记录ID
                                        recordId.set(recordIdStr);
                                    }
                                    LkeReplyResponse lkeReplyResponse = new LkeReplyResponse(
                                            chatMessageId,
                                            sessionId,
                                            request.getModel(),
                                            ChatRoleType.AI,
                                            replyContent,
                                            MessageType.MARKDOWN,
                                            isFinal);
                                    sink.next(StreamData.data(new AIStreamMessageVo(lkeReplyResponse)));
                                    sb.append(replyContent);
                                    if (isFinal) {
                                        // 储存消息
                                        AiChatMessage message = new AiChatMessage(lkeReplyResponse);
                                        message.setId(chatMessageId);
                                        message.setContent(sb.toString());
                                        message.setRecordId(recordId.get());
                                        aiChatMessageService.saveMessageAsync(message);
                                        try {
                                            for (Future future : futures) {
                                                if (future != null) {
                                                    future.get(); // 等待总结任务完成
                                                }
                                            }
                                        } catch (InterruptedException | ExecutionException e) {
                                            log.error("总结任务执行失败", e);
                                        }
                                        log.info("AI响应流完成");
                                        Integer spendTime = (int) Duration.between(replyStartTime, LocalDateTime.now()).toMillis();
                                        socket.off("reply");
                                        sink.next(StreamData.done(new AISteamDoneVo(replyStartTime.toString(), spendTime)));
                                        sink.complete();
                                    }
                                } catch (JSONException e) {
                                    log.error("Error processing reply: {}", e.getMessage());
                                    socket.off("reply");
                                    if (recordId.get() != null) {
                                        // 停止生成
                                        JSONObject payload = new JSONObject();
                                        try {
                                            payload.put("record_id", recordId.get());
                                        } catch (JSONException ex) {
                                            throw new RuntimeException(ex);
                                        }
                                        log.info("payload: {}", payload);
                                        socket.emit("stop_generation", payload);
                                    }
                                    sink.error(e);
                                }
                            }
                        })
                        // 监听error事件
                        .once("error", args -> {
                            socket.off("reply");
                            sink.error(new RuntimeException("Socket error: " + Arrays.toString(args)));
                        });
            });
        } catch (InterruptedException | JSONException e) {
            throw new RuntimeException(e);
        } finally {
            REQUEST.release();
        }
    }


    /**
     * 异步处理总结问题
     *
     * @param sink
     * @param chatSession
     * @param content
     * @return
     */
    private Future<Void> summaryQuestionAsync(FluxSink<StreamData<?>> sink, AiChatSession chatSession, String content) {
        return taskExecutor.submit(() -> {
            String textContent = MarkdownUtil.extractText(content);
            String summary = aiSummaryService.getSummaryFromText(textContent);
            chatSession.setSessionName(summary);
            aiChatSessionService.save(chatSession);
            sink.next(StreamData.updateSession(new AIChatSessionVo(chatSession)));
            return null;
        });
    }

    private static AiChatSession createAiChatSession(Long sessionId, Long userId) {
        AiChatSession chatSession = new AiChatSession();
        chatSession.setId(sessionId);
        chatSession.setSessionName("New Chat");
        chatSession.setUserId(userId);
        chatSession.setCreateTime(LocalDateTime.now());
        chatSession.setUpdateTime(LocalDateTime.now());
        return chatSession;
    }
}
